import os
import torch

class Config:
    # 数据路径配置 - 需要用户设置
    DATA_ROOT = r"E:\image_data"  # 替换为您的数据根目录

    # 训练参数
    BATCH_SIZE = 8
    NUM_EPOCHS = 80
    LEARNING_RATE = 3e-4
    ALPHA = 0.2  # 损失函数参数
    MBT_KERNELS = [5, 9]  # MBT增强核大小

    # 设备配置
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

    # 输入尺寸
    IMAGE_SIZE = (256, 256)

    # 模型保存路径
    SAVE_DIR = "saved_models"
    os.makedirs(SAVE_DIR, exist_ok=True)

    # 数据划分比例 (训练:验证:测试)
    TRAIN_RATIO = 0.8
    VAL_RATIO = 0.1
    TEST_RATIO = 0.1

    # 随机种子
    RANDOM_SEED = 42

    PATIENCE = 40

def get_sample_paths():
    """获取所有样本路径 (PNG和JSON文件对)"""
    sample_paths = []
    data_root = Config.DATA_ROOT

    # 遍历所有样本文件夹
    for folder in os.listdir(data_root):
        folder_path = os.path.join(data_root, folder)
        if not os.path.isdir(folder_path):
            continue

        # 检查文件夹中的文件
        png_files = [f for f in os.listdir(folder_path) if f.endswith(".png")]
        json_files = [f for f in os.listdir(folder_path) if f.endswith(".json")]

        # 匹配PNG和JSON文件
        for png_file in png_files:
            base_name = os.path.splitext(png_file)[0]
            json_file = f"{base_name}.json"

            if json_file in json_files:
                png_path = os.path.join(folder_path, png_file)
                json_path = os.path.join(folder_path, json_file)
                sample_paths.append((png_path, json_path))

    print(f"Found {len(sample_paths)} valid samples in {data_root}")
    return sample_paths

